import cv2 as cv
import numpy as np


def pad_img(img):
    H, W = img.shape[:2]
    if H > W:
        padding = H - W
        img = np.concatenate([img, np.zeros([H, padding, 3], dtype=np.uint8)], axis=1)
    else:
        padding = W - H
        img = np.concatenate([img, np.zeros([padding, W, 3], dtype=np.uint8)], axis=0)
    return img


def GaussianPyramid(img, level, name='img', is_show=False):
    MAX_LEVEL = 9
    SIZE = 2 ** MAX_LEVEL
    img = pad_img(img)
    img = cv.resize(img, (SIZE, SIZE), interpolation=cv.INTER_CUBIC)

    if is_show:
        cv.imshow(f'{name}_#0', img)

    pyramid_arr = []
    LEVELS = MAX_LEVEL - 5
    if level > LEVELS:
        level = LEVELS
    tmp = img
    for i_level in range(level):
        target = cv.pyrDown(tmp)
        if is_show:
            cv.imshow(f'{name}_gau#{i_level + 1}', target)
        pyramid_arr.append(target)
        tmp = target

    pyramid_arr.insert(0, img)

    return pyramid_arr


def LaplacianPyramid(img, level, name='img', is_show=False):
    gaussian_arr = GaussianPyramid(img, level, name)
    arr_len = len(gaussian_arr)
    for i_level in range(arr_len - 1):
        a = gaussian_arr[i_level]
        b = gaussian_arr[i_level + 1]
        b = cv.pyrUp(b)
        target = cv.subtract(a, b)
        if is_show:
            cv.imshow(f'{name}_lap#{i_level + 1}', target)
        gaussian_arr.append(target)
    return gaussian_arr


if '__main__' == __name__:
    img_path = '../../../large_data/pic/dog_bird.jpg'
    img = cv.imread(img_path, cv.IMREAD_COLOR)

    pyramid_arr = GaussianPyramid(img, 3, 'dog_bird', True)

    laplacian_arr = LaplacianPyramid(img, 3, 'dog_bird', True)

    cv.waitKey(0)
    cv.destroyAllWindows()
    print('Over')
